import time
from contextlib import contextmanager
from typing import Generator

from sqlalchemy import create_engine, event, Engine
from sqlalchemy.engine.base import Connection  # 显式导入 Connection 类型
from sqlalchemy.orm import sessionmaker, Session

from api.fm.configs import ConfigManager
from api.fm.models.base import Base


# 定义事件监听器
@event.listens_for(Session, "before_commit")
def before_commit(session):
    # print("Before commit")
    pass


@event.listens_for(Session, "before_flush")
def before_flush(session, flush_context, instances):
    # print("Before flush")
    pass


@event.listens_for(Session, "after_flush")
def after_flush(session, flush_context):
    # print("After flush")
    pass


@event.listens_for(Session, "after_commit")
def after_commit(session):
    # print("After commit")
    pass


@event.listens_for(Session, "after_rollback")
def after_rollback(session):
    # print("After rollback")
    pass


class DatabaseManager:
    def __init__(self, config_manager: 'ConfigManager'):
        # 获取配置
        self.config = config_manager.get_config()

        # 创建数据库引擎
        database_uri = self.config.get("SQLALCHEMY_DATABASE_URI")
        engine_options = self.config.get("SQLALCHEMY_ENGINE_OPTIONS", {})

        if not database_uri:
            raise ValueError("SQLALCHEMY_DATABASE_URI is required in the config.")

        self.engine: Engine = create_engine(
            database_uri,
            **engine_options
        )

        # 创建会话工厂
        self.SessionLocal = sessionmaker(
            autocommit=False,
            autoflush=False,
            bind=self.engine
        )

    @contextmanager
    def get_db_session(self) -> Generator[Session, None, None]:
        """
        获取数据库会话上下文管理器，自动关闭连接并打印执行时间
        """
        start_time = time.time()
        db: Session = self.SessionLocal()  # 创建数据库会话
        try:
            yield db  # 将会话传递给调用者
        except Exception as e:
            db.rollback()  # 如果发生异常，回滚事务
            raise e  # 重新抛出异常
        finally:
            end_time = time.time()
            print(f"Database execution time: {end_time - start_time:.6f} seconds")
            db.close()  # 最终关闭会话

    @contextmanager
    def get_db_connection(self) -> Generator[Connection, None, None]:
        """
        获取数据库连接上下文管理器，自动关闭连接
        """
        connection: Connection = self.engine.connect()
        try:
            yield connection  # 将连接传递给调用者
        finally:
            connection.close()  # 最终关闭连接

    def create_all(self):
        """
        创建所有数据库表
        """
        Base.metadata.create_all(self.engine)
        print("All tables created.")
